package org.Yiran.timetale_re.datagen;

import net.minecraft.data.loot.BlockLootSubProvider;
import net.minecraft.world.flag.FeatureFlags;
import net.minecraft.world.item.Item;
import net.minecraft.world.item.Items;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.state.properties.IntegerProperty;
import net.minecraft.world.level.storage.loot.LootPool;
import net.minecraft.world.level.storage.loot.LootTable;
import net.minecraft.world.level.storage.loot.entries.LootItem;
import net.minecraft.world.level.storage.loot.functions.SetItemCountFunction;
import net.minecraft.world.level.storage.loot.predicates.LootItemBlockStatePropertyCondition;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
import net.minecraft.world.level.storage.loot.providers.number.ConstantValue;
import net.minecraft.world.level.storage.loot.providers.number.UniformGenerator;
import net.minecraftforge.registries.RegistryObject;
import org.Yiran.timetale_re.block.BlocksRegister;
import org.Yiran.timetale_re.block.crystal.base.CrystalBaseBlock;
import org.Yiran.timetale_re.item.ItemsRegister;
import org.Yiran.timetale_re.item.MaterialRegister;

import java.util.*;

import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class ModBlockLootTables extends BlockLootSubProvider {
    private static final Logger LOGGER = LoggerFactory.getLogger(ModBlockLootTables.class);

    public ModBlockLootTables() {
        super(Set.of(), FeatureFlags.REGISTRY.allFlags());
    }

    @Override
    protected void generate() {
        LOGGER.info("开始生成方块战利品表...");
        
        // 为长桌生成战利品表
        LOGGER.info("正在生成长桌战利品表...");
        add(BlocksRegister.LONG_TABLE.get(), createSingleItemTable(BlocksRegister.LONG_TABLE.get()));
        
        // 为椅子生成战利品表
        LOGGER.info("正在生成椅子战利品表...");
        add(BlocksRegister.CHAIR.get(), createSingleItemTable(BlocksRegister.CHAIR.get()));
        
        // 为所有注册的水晶方块生成战利品表
        LOGGER.info("正在生成水晶方块战利品表...");
        generateCrystalLootTables();
        
        // 为材料方块生成战利品表
        LOGGER.info("正在生成材料方块战利品表...");
        generateMaterialBlockLootTables();
        
        LOGGER.info("方块战利品表生成完成！");
    }

    /**
     * 为水晶方块生成战利品表
     * 根据方块的level属性决定掉落物数量
     */
    private void generateCrystalLootTables() {
        // 获取所有水晶方块
        List<Block> crystalBlocks = getCrystalBlocks();
        
        LOGGER.info("检测到 {} 个水晶方块，开始生成战利品表...", crystalBlocks.size());
        
        // 为所有水晶方块生成战利品表
        crystalBlocks.forEach(this::processCrystalBlock);
        
        LOGGER.info("水晶方块战利品表生成完成！");
    }
    
    /**
     * 获取所有注册的水晶方块
     * 
     * @return 水晶方块列表
     */
    private List<Block> getCrystalBlocks() {
        return BlocksRegister.BLOCKS.getEntries().stream()
                .map(RegistryObject::get)
                .filter(block -> block instanceof CrystalBaseBlock)
                .toList();
    }
    
    /**
     * 处理单个水晶方块的战利品表生成
     * 
     * @param block 水晶方块
     */
    private void processCrystalBlock(Block block) {
        String blockName = getBlockName(block);
        String itemName = getItemNameFromBlockName(blockName);
        
        LOGGER.debug("正在为方块 {} 生成战利品表，将使用物品 {}", blockName, itemName);
        
        // 特殊处理energy_crystal_cluster
        if (blockName.equals("energy_crystal_cluster")) {
            handleEnergyCrystalCluster(block);
        } else {
            // 为其他水晶方块生成标准战利品表
            handleStandardCrystalBlock(block, itemName);
        }
    }
    
    /**
     * 获取方块名称（去除命名空间前缀）
     * 
     * @param block 方块
     * @return 方块名称
     */
    private String getBlockName(Block block) {
        return block.getDescriptionId().replace("block.timetale_re.", "");
    }
    
    /**
     * 根据方块名称推导出对应的物品名称
     * 规则是：
     * 1. 将"_crystal"替换为"_gem"
     * 2. 将"cluster"替换为"gem"
     * 
     * @param blockName 方块名称
     * @return 物品名称
     */
    private String getItemNameFromBlockName(String blockName) {
        return blockName.replace("_crystal", "_gem").replace("cluster", "gem");
    }
    
    /**
     * 处理能晶簇方块的特殊战利品表生成
     * 
     * @param block 能晶簇方块
     */
    private void handleEnergyCrystalCluster(Block block) {
        String blockName = getBlockName(block);
        LOGGER.debug("方块 {} 是能晶簇，使用特殊战利品表生成逻辑", blockName);
        add(block, createEnergyCrystalClusterDrops(block));
    }
    
    /**
     * 处理标准水晶方块的战利品表生成
     * 
     * @param block 水晶方块
     * @param itemName 对应的物品名称
     */
    private void handleStandardCrystalBlock(Block block, String itemName) {
        String blockName = getBlockName(block);
        LOGGER.debug("方块 {} 使用标准战利品表生成逻辑", blockName);
        add(block, createCrystalDrops(block, itemName));
    }

    /**
     * 为材料相关方块生成战利品表
     */
    private void generateMaterialBlockLootTables() {
        LOGGER.info("开始为 {} 种材料生成方块战利品表...", MaterialRegister.getRegisteredMaterials().size());
        
        // 统计各类方块数量
        int[] counts = processMaterialBlocks();
        int oreCount = counts[0];
        int blockCount = counts[1];
        int rawBlockCount = counts[2];
        
        LOGGER.info("材料方块战利品表生成完成！矿石方块: {} 个, 存储块: {} 个, 粗矿块: {} 个", oreCount, blockCount, rawBlockCount);
    }
    
    /**
     * 处理所有材料的方块战利品表生成
     * 
     * @return 包含[矿石数量, 存储块数量, 粗矿块数量]的数组
     */
    private int[] processMaterialBlocks() {
        int oreCount = 0;
        int blockCount = 0;
        int rawBlockCount = 0;
        
        for (MaterialRegister.MaterialDefinition material : MaterialRegister.getRegisteredMaterials()) {
            // 处理矿石方块
            oreCount += processOreBlocks(material);
            
            // 处理存储块
            blockCount += processStorageBlocks(material);
            
            // 处理粗矿块
            rawBlockCount += processRawBlocks(material);
        }
        
        return new int[]{oreCount, blockCount, rawBlockCount};
    }
    
    /**
     * 处理材料的矿石方块（普通矿石和深板岩矿石）
     * 
     * @param material 材料定义
     * @return 处理的矿石方块数量
     */
    private int processOreBlocks(MaterialRegister.MaterialDefinition material) {
        int count = 0;
        
        // 矿石方块
        if (material.hasOre()) {
            if (material.getOre() != null) {
                LOGGER.debug("为材料 {} 生成普通矿石战利品表", material.getName());
                // 使用createOreDrop方法替代createSingleItemTable以支持时运附魔
                add(material.getOre().get(), createOreDrop(material.getOre().get(), material.getRaw().get()));
                count++;
            }
            if (material.getDeepslateOre() != null) {
                LOGGER.debug("为材料 {} 生成深板岩矿石战利品表", material.getName());
                // 使用createOreDrop方法替代createSingleItemTable以支持时运附魔
                add(material.getDeepslateOre().get(), createOreDrop(material.getDeepslateOre().get(), material.getRaw().get()));
                count++;
            }
        }
        
        return count;
    }
    
    /**
     * 处理材料的存储块
     * 
     * @param material 材料定义
     * @return 处理的存储块数量（0或1）
     */
    private int processStorageBlocks(MaterialRegister.MaterialDefinition material) {
        int count = 0;
        
        // 存储块
        if (material.hasType(MaterialRegister.MaterialType.BLOCK) && material.getBlock() != null) {
            LOGGER.debug("为材料 {} 生成存储块战利品表", material.getName());
            add(material.getBlock().get(), createSingleItemTable(material.getBlockItem().get()));
            count++;
        }
        
        return count;
    }
    
    /**
     * 处理材料的粗矿块
     * 
     * @param material 材料定义
     * @return 处理的粗矿块数量（0或1）
     */
    private int processRawBlocks(MaterialRegister.MaterialDefinition material) {
        int count = 0;
        
        // 粗矿块
        if (material.hasRawBlock() && material.getRawBlock() != null) {
            LOGGER.debug("为材料 {} 生成粗矿块战利品表", material.getName());
            add(material.getRawBlock().get(), createSingleItemTable(material.getRawBlockItem().get()));
            count++;
        }
        
        return count;
    }

    /**
     * 创建水晶方块的战利品表
     * 根据level属性掉落不同数量的宝石
     */
    protected LootTable.Builder createCrystalDrops(Block block, String itemName) {
        // 获取方块的等级属性
        IntegerProperty levelProperty = ((CrystalBaseBlock) block).getLevelProperty();
        
        // 创建等级条件
        Map<Integer, LootItemCondition.Builder> levelConditions = createLevelConditions(block, levelProperty, 1, 2, 3);
        LootItemCondition.Builder level1Condition = levelConditions.get(1);
        LootItemCondition.Builder level2Condition = levelConditions.get(2);
        LootItemCondition.Builder level3Condition = levelConditions.get(3);

        // 根据物品名获取物品
        Item item = getItemFromName(itemName);
        
        return LootTable.lootTable()
                .withPool(applyExplosionCondition(block, LootPool.lootPool()
                        .setRolls(ConstantValue.exactly(1))
                        .add(LootItem.lootTableItem(item)
                                .when(level1Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(1))))
                        .add(LootItem.lootTableItem(item)
                                .when(level2Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(2))))
                        .add(LootItem.lootTableItem(item)
                                .when(level3Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(3))))));
    }

    /**
     * 创建能晶簇的特殊战利品表
     */
    protected LootTable.Builder createEnergyCrystalClusterDrops(Block block) {
        // 获取方块的等级属性
        IntegerProperty levelProperty = getLevelPropertyForBlock(block);
        
        // 创建等级条件
        Map<Integer, LootItemCondition.Builder> levelConditions = createLevelConditions(block, levelProperty, 1, 2, 3, 4);
        LootItemCondition.Builder level1Condition = levelConditions.get(1);
        LootItemCondition.Builder level2Condition = levelConditions.get(2);
        LootItemCondition.Builder level3Condition = levelConditions.get(3);
        LootItemCondition.Builder level4Condition = levelConditions.get(4);

        return LootTable.lootTable()
                .withPool(applyExplosionCondition(block, LootPool.lootPool()
                        .setRolls(ConstantValue.exactly(1))
                        .add(LootItem.lootTableItem(Items.AIR)
                                .when(level1Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(0))))
                        .add(LootItem.lootTableItem(Items.AIR)
                                .when(level2Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(0))))
                        .add(LootItem.lootTableItem(Items.AIR)
                                .when(level3Condition)
                                .apply(SetItemCountFunction.setCount(ConstantValue.exactly(0))))
                        .add(LootItem.lootTableItem(getItemFromName("energy_gem"))
                                .when(level4Condition)
                                .apply(SetItemCountFunction.setCount(UniformGenerator.between(2.0F, 2.0F))))));
    }
    
    /**
     * 获取方块的等级属性
     * 
     * @param block 方块
     * @return 等级属性
     */
    private IntegerProperty getLevelPropertyForBlock(Block block) {
        if (block instanceof CrystalBaseBlock crystalBaseBlock) {
            return crystalBaseBlock.getLevelProperty();
        }
        // 默认返回空属性，但这不应该发生
        throw new IllegalArgumentException("Block is not a CrystalBaseBlock: " + block);
    }

    /**
     * 创建指定等级的方块状态条件
     * @param block 方块
     * @param levelProperty 等级属性
     * @param levels 需要创建条件的等级
     * @return 包含各等级条件的映射
     */
    private Map<Integer, LootItemCondition.Builder> createLevelConditions(Block block, IntegerProperty levelProperty, Integer... levels) {
        Map<Integer, LootItemCondition.Builder> conditions = new HashMap<>();
        for (int level : levels) {
            conditions.put(level, LootItemBlockStatePropertyCondition.hasBlockStateProperties(block)
                    .setProperties(net.minecraft.advancements.critereon.StatePropertiesPredicate.Builder.properties()
                            .hasProperty(levelProperty, level)));
        }
        return conditions;
    }

    /**
     * 根据物品名获取物品实例
     * @param itemName 物品名
     * @return 对应的物品实例
     */
    private Item getItemFromName(String itemName) {
        return switch (itemName) {
            case "time_gem" -> ItemsRegister.TIME_GEM.get();
            case "dimension_gem" -> ItemsRegister.DIMENSION_GEM.get();
            case "space_gem" -> ItemsRegister.SPACE_GEM.get();
            case "chaos_gem" -> ItemsRegister.CHAOS_GEM.get();
            case "energy_gem" -> ItemsRegister.ENERGY_GEM.get();
            default -> Items.AIR;
        };
    }

    @Override
    protected @NotNull Iterable<Block> getKnownBlocks() {
        List<Block> blocks = new ArrayList<>();
        
        // 添加所有注册的方块
        BlocksRegister.BLOCKS.getEntries().stream()
                .map(RegistryObject::get)
                .forEach(blocks::add);
                
        // 添加所有材料方块
        for (MaterialRegister.MaterialDefinition material : MaterialRegister.getRegisteredMaterials()) {
            if (material.hasOre()) {
                if (material.getOre() != null) {
                    blocks.add(material.getOre().get());
                }
                if (material.getDeepslateOre() != null) {
                    blocks.add(material.getDeepslateOre().get());
                }
            }
            
            if (material.hasType(MaterialRegister.MaterialType.BLOCK) && material.getBlock() != null) {
                blocks.add(material.getBlock().get());
            }
            
            if (material.hasRawBlock() && material.getRawBlock() != null) {
                blocks.add(material.getRawBlock().get());
            }
        }
        
        return blocks;
    }
}